import os.path
import random

import tensorflow as tf
import librosa
import numpy as np
from tqdm import tqdm

def _float_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def _int64_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))

def data_example(data, label):
    feature = {
        'data': _float_feature(data),
        'label': _int64_feature(label),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))

def create_data_tfrecord(data_list_path, save_path, label_map):
    data = []
    clsPath = os.listdir(data_list_path)
    for single in clsPath:
        singelePath = os.path.join(data_list_path, single)
        singelePathdir = os.listdir(singelePath)
        for s in singelePathdir:
            spath = os.path.join(singelePath, s)
            label = label_map[single]
            data.append([spath, label])

    with tf.io.TFRecordWriter(save_path) as writer:
        for d in tqdm(data):
            try:
                path, label = d[0], d[1]
                wav, sr = librosa.load(path, sr=16000)
                intervals = librosa.effects.split(wav, top_db=20)
                wav_output = []
                wav_len = int(16000 * 2.04)
                for sliced in intervals:
                    wav_output.extend(wav[sliced[0]:sliced[1]])
                for i in range(5):
                    if len(wav_output) > wav_len:
                        l = len(wav_output) - wav_len
                        r = random.randint(0, l)
                        wav_output = wav_output[r:wav_len + r]
                    else:
                        wav_output.extend(np.zeros(shape=[wav_len - len(wav_output)], dtype=np.float32))
                    wav_output = np.array(wav_output)
                    ps = librosa.feature.melspectrogram(y=wav_output, sr=sr, hop_length=256).reshape(-1).tolist()
                    if len(ps) != 128 * 128: continue
                    tf_example = data_example(ps, int(label))
                    writer.write(tf_example.SerializeToString())
                    if len(wav_output) <= wav_len:
                        break
            except Exception as e:
                print(e)

if __name__ == '__main__':
    # 定义类别标签映射
    label_map = {
        'abnormal_treble': 0,
        'normal_bass': 1,
        'normal_bass_inside_the_arriage': 2,
        'normal_treble': 3,
        'normal_treble_inside_the_arriage': 4,
        'resonant_bass_inside_the_carriage': 5,
        'resonant_treble_inside_the_carriage': 6
    }

    # 创建TFRecord文件
    create_data_tfrecord('D:/Pycharm/trainvoice5/train', 'D:/Pycharm/ResnetData/train.tfrecord', label_map)
    create_data_tfrecord('D:/Pycharm/trainvoice5/test', 'D:/Pycharm/ResnetData/test.tfrecord', label_map)
